from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage
import json

class StateManagementState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    user_info: Annotated[Dict[str, Any], "用户信息"]
    session_data: Annotated[Dict[str, Any], "会话数据"]
    step_count: Annotated[int, "步骤计数"]
    history: Annotated[list, "操作历史"]

class StateManagementDemo:
    """状态管理演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建状态管理图"""
        
        workflow = StateGraph(StateManagementState)
        
        # 添加节点
        workflow.add_node("initialize_session", self._initialize_session)
        workflow.add_node("process_user_input", self._process_user_input)
        workflow.add_node("update_user_info", self._update_user_info)
        workflow.add_node("track_activity", self._track_activity)
        workflow.add_node("generate_response", self._generate_response)
        workflow.add_node("save_session", self._save_session)
        
        # 设置入口点
        workflow.set_entry_point("initialize_session")
        
        # 设置边
        workflow.add_edge("initialize_session", "process_user_input")
        workflow.add_edge("process_user_input", "update_user_info")
        workflow.add_edge("update_user_info", "track_activity")
        workflow.add_edge("track_activity", "generate_response")
        workflow.add_edge("generate_response", "save_session")
        workflow.add_edge("save_session", END)
        
        return workflow.compile()
    
    def _initialize_session(self, state: StateManagementState) -> StateManagementState:
        """初始化会话状态"""
        messages = state.get("messages", [])
        if not messages:
            messages = [HumanMessage(content="开始新会话")]
        
        return {
            "messages": messages,
            "result": "会话已初始化",
            "user_info": {
                "session_id": f"session_{hash(str(messages))}",
                "start_time": "2024-01-01 10:00:00",
                "preferences": {}
            },
            "session_data": {
                "conversation_count": 0,
                "last_activity": "初始化",
                "context": {}
            },
            "step_count": 1,
            "history": ["会话初始化"]
        }
    
    def _process_user_input(self, state: StateManagementState) -> StateManagementState:
        """处理用户输入"""
        messages = state["messages"]
        content = messages[-1].content if messages else "默认消息"
        
        # 更新会话数据
        session_data = state["session_data"]
        session_data["conversation_count"] += 1
        session_data["last_activity"] = "处理用户输入"
        
        # 更新历史
        history = state["history"] + [f"处理输入: {content}"]
        
        return {
            "messages": messages,
            "result": f"已处理用户输入: {content}",
            "user_info": state["user_info"],
            "session_data": session_data,
            "step_count": state["step_count"] + 1,
            "history": history
        }
    
    def _update_user_info(self, state: StateManagementState) -> StateManagementState:
        """更新用户信息"""
        user_info = state["user_info"]
        content = state["messages"][-1].content if state["messages"] else ""
        
        # 根据输入内容更新用户偏好
        if "喜欢" in content:
            user_info["preferences"]["likes"] = content
        elif "不喜欢" in content:
            user_info["preferences"]["dislikes"] = content
        
        # 更新历史
        history = state["history"] + ["更新用户信息"]
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "user_info": user_info,
            "session_data": state["session_data"],
            "step_count": state["step_count"] + 1,
            "history": history
        }
    
    def _track_activity(self, state: StateManagementState) -> StateManagementState:
        """跟踪活动"""
        session_data = state["session_data"]
        session_data["last_activity"] = "活动跟踪"
        session_data["context"]["current_step"] = state["step_count"]
        
        # 更新历史
        history = state["history"] + [f"跟踪活动: 步骤{state['step_count']}"]
        
        return {
            "messages": state["messages"],
            "result": state["result"],
            "user_info": state["user_info"],
            "session_data": session_data,
            "step_count": state["step_count"] + 1,
            "history": history
        }
    
    def _generate_response(self, state: StateManagementState) -> StateManagementState:
        """生成响应"""
        user_info = state["user_info"]
        session_data = state["session_data"]
        
        # 基于状态生成个性化响应
        response = f"基于您的偏好({user_info['preferences']})和会话历史({session_data['conversation_count']}次对话)，我为您生成了个性化响应。"
        
        # 更新历史
        history = state["history"] + ["生成响应"]
        
        return {
            "messages": state["messages"] + [AIMessage(content=response)],
            "result": response,
            "user_info": user_info,
            "session_data": session_data,
            "step_count": state["step_count"] + 1,
            "history": history
        }
    
    def _save_session(self, state: StateManagementState) -> StateManagementState:
        """保存会话状态"""
        # 模拟保存到数据库
        session_summary = {
            "session_id": state["user_info"]["session_id"],
            "total_steps": state["step_count"],
            "conversation_count": state["session_data"]["conversation_count"],
            "user_preferences": state["user_info"]["preferences"],
            "execution_history": state["history"]
        }
        
        # 更新历史
        history = state["history"] + ["保存会话状态"]
        
        return {
            "messages": state["messages"],
            "result": f"会话已保存: {session_summary}",
            "user_info": state["user_info"],
            "session_data": state["session_data"],
            "step_count": state["step_count"],
            "history": history
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行状态管理图"""
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "开始新会话"))],
            "result": "",
            "user_info": {},
            "session_data": {},
            "step_count": 0,
            "history": []
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "session_summary": {
                "session_id": result["user_info"]["session_id"],
                "total_steps": result["step_count"],
                "user_preferences": result["user_info"]["preferences"],
                "execution_history": result["history"]
            },
            "final_result": result["result"]
        }
